import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dropout, Dense, SimpleRNN
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error
import math
import os
import shutil
import time
import subprocess
import sys
# 获取当前脚本所在的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
print(f"脚本目录: {script_dir}")
print(f"当前工作目录: {os.getcwd()}")

# 强制使用CPU
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# 加载数据
data_path = os.path.join(script_dir, 'SH600519.csv')
print(f"数据路径: {data_path}")
maotai = pd.read_csv(data_path)
training_set = maotai.iloc[0:2426 - 300, 2:3].values
test_set = maotai.iloc[2426 - 300:, 2:3].values

# 归一化
sc = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
test_set_scaled = sc.transform(test_set)

# 定义不同的参数组合
param_combinations = [
    {'name': 'Baseline', 'neurons1': 80, 'neurons2': 100, 'dropout': 0.2, 'window': 60, 'batch_size': 64, 'epochs': 50},
    {'name': 'Complex Model', 'neurons1': 128, 'neurons2': 256, 'dropout': 0.3, 'window': 60, 'batch_size': 64,
     'epochs': 70},
    {'name': 'Simple Model', 'neurons1': 32, 'neurons2': 64, 'dropout': 0.1, 'window': 60, 'batch_size': 32,
     'epochs': 30},
    {'name': 'Long Window', 'neurons1': 80, 'neurons2': 100, 'dropout': 0.2, 'window': 90, 'batch_size': 64,
     'epochs': 50},
    {'name': 'Short Window', 'neurons1': 80, 'neurons2': 100, 'dropout': 0.2, 'window': 30, 'batch_size': 64,
     'epochs': 50},
    {'name': 'High Dropout', 'neurons1': 80, 'neurons2': 100, 'dropout': 0.5, 'window': 60, 'batch_size': 64,
     'epochs': 50}
]

# 存储所有结果
all_results = []
all_predictions = []

# 使用系统字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统
plt.rcParams['axes.unicode_minus'] = False

# 设置对比图颜色 - 为每个实验组名称分配固定颜色
name_to_color = {
    'Baseline': 'b',         # 蓝色
    'Complex Model': 'g',    # 绿色
    'Simple Model': 'r',     # 红色
    'Long Window': 'c',      # 青色
    'Short Window': 'm',     # 品红
    'High Dropout': 'y'      # 黄色
}
# 扩展线型列表
line_styles = [
    '-',  # 实线
    '--',  # 虚线
    '-.',  # 点划线
    ':',  # 点线
    (0, (5, 5)),  # 稀疏虚线 (自定义)
    (0, (3, 1, 1, 1)),  # 点虚线 (自定义)
    (0, (1, 1)),  # 密集虚线 (自定义)
    (0, (5, 1)),  # 长虚线 (自定义)
]

# 循环执行不同参数组合
for idx, params in enumerate(param_combinations):
    print(f"\n{'=' * 50}")
    print(f"开始实验 #{idx + 1}/{len(param_combinations)}: {params['name']}")
    # ... 中间代码保持不变 ...

    # 为当前组合分配颜色和线型
    color = name_to_color.get(params['name'], 'k')  # 默认黑色
    line_style = line_styles[idx % len(line_styles)]

# 循环执行不同参数组合
for idx, params in enumerate(param_combinations):
    print(f"\n{'=' * 50}")
    print(f"开始实验 #{idx + 1}/{len(param_combinations)}: {params['name']}")
    print(f"参数配置: neurons1={params['neurons1']}, neurons2={params['neurons2']}, "
          f"dropout={params['dropout']}, window={params['window']}, "
          f"batch_size={params['batch_size']}, epochs={params['epochs']}")
    print(f"{'=' * 50}")

    start_time = time.time()



    # 根据当前参数构建训练数据
    x_train, y_train = [], []
    for i in range(params['window'], len(training_set_scaled)):
        x_train.append(training_set_scaled[i - params['window']:i, 0])
        y_train.append(training_set_scaled[i, 0])

    # 打乱数据
    np.random.seed(7)
    indices = np.arange(len(x_train))
    np.random.shuffle(indices)
    x_train = np.array(x_train)[indices]
    y_train = np.array(y_train)[indices]
    x_train = np.reshape(x_train, (x_train.shape[0], params['window'], 1))

    # 构建测试数据
    x_test, y_test = [], []
    for i in range(params['window'], len(test_set_scaled)):
        x_test.append(test_set_scaled[i - params['window']:i, 0])
        y_test.append(test_set_scaled[i, 0])
    x_test, y_test = np.array(x_test), np.array(y_test)
    x_test = np.reshape(x_test, (x_test.shape[0], params['window'], 1))

    # 为当前参数组合计算真实股价（从当前窗口开始）
    current_real_stock_price = sc.inverse_transform(test_set_scaled[params['window']:])

    # 模型定义
    model = tf.keras.Sequential([
        SimpleRNN(params['neurons1'], return_sequences=True),
        Dropout(params['dropout']),
        SimpleRNN(params['neurons2']),
        Dropout(params['dropout']),
        Dense(1)
    ])

    model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
                  loss='mean_squared_error')

    # 为每个参数组合创建单独的检查点目录
    checkpoint_dir = os.path.join(script_dir, f"checkpoints_{params['name']}")
    checkpoint_save_path = os.path.join(checkpoint_dir, "rnn_stock.ckpt")

    # 确保文件目录存在
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir, exist_ok=True)

    # 检查点回调
    cp_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath=checkpoint_save_path,
        save_weights_only=True,
        save_best_only=True,
        monitor='val_loss'
    )

    # 模型训练
    print("开始训练模型...")
    try:
        history = model.fit(
            x_train, y_train,
            batch_size=params['batch_size'],
            epochs=params['epochs'],
            validation_data=(x_test, y_test),
            validation_freq=1,
            callbacks=[cp_callback],
            verbose=1
        )
    except Exception as e:
        print(f"训练过程中发生错误: {e}")
        print("尝试不使用回调函数训练...")
        history = model.fit(
            x_train, y_train,
            batch_size=params['batch_size'],
            epochs=params['epochs'],
            validation_data=(x_test, y_test),
            validation_freq=1,
            verbose=1
        )

    # 预测和评估
    predicted_stock_price = model.predict(x_test)
    predicted_stock_price = sc.inverse_transform(predicted_stock_price)

    # 计算评估指标
    mse = mean_squared_error(current_real_stock_price, predicted_stock_price)
    rmse = math.sqrt(mse)
    mae = mean_absolute_error(current_real_stock_price, predicted_stock_price)

    # 保存结果
    # 保存结果
    result = {
        'name': params['name'],
        'params': params,
        'mse': mse,
        'rmse': rmse,
        'mae': mae,
        'predicted': predicted_stock_price,
        'real': current_real_stock_price,
        'window': params['window'],
        'training_time': time.time() - start_time,
        'color': color,  # 使用名称映射的颜色
        'line_style': line_style,  # 基于索引的线型
        'history': history
    }

    all_results.append(result)
    all_predictions.append(predicted_stock_price)

    print(f"\n实验 '{params['name']}' 完成!")
    print(f'测试样本数量: {len(current_real_stock_price)}')
    print(f'预测样本数量: {len(predicted_stock_price)}')
    print(f'均方误差(MSE): {mse:.6f}')
    print(f'均方根误差(RMSE): {rmse:.6f}')
    print(f'平均绝对误差(MAE): {mae:.6f}')
    print(f'训练时间: {result["training_time"]:.2f}秒')

    # 清理检查点目录
    if os.path.exists(checkpoint_dir):
        print(f"清理临时目录: {checkpoint_dir}")
        shutil.rmtree(checkpoint_dir)

# 结果可视化
# 1. 绘制所有参数组合的预测结果对比
plt.figure(figsize=(16, 10))

# 绘制真实股价（使用第一个参数组合的窗口作为参考）
ref_window = param_combinations[0]['window']
plt.plot(sc.inverse_transform(test_set_scaled[ref_window:]),
         color='k', linewidth=3, label='真实股价', linestyle='-')

for result in all_results:
    # 对齐不同窗口大小的预测结果
    offset = result['window'] - ref_window
    if offset < 0:
        # 如果窗口更小，截断预测结果以对齐
        adjusted_pred = result['predicted'][-offset:]
    elif offset > 0:
        # 如果窗口更大，在前面填充NaN
        adjusted_pred = np.concatenate([np.full((offset, 1), np.nan), result['predicted']], axis=0)
    else:
        adjusted_pred = result['predicted']

    plt.plot(adjusted_pred,
             color=result['color'],
             linestyle=result['line_style'],
             linewidth=2,
             label=f"{result['name']} (RMSE:{result['rmse']:.2f})")

plt.title('不同参数组合的股价预测效果对比', fontsize=18)
plt.xlabel('时间', fontsize=14)
plt.ylabel('价格', fontsize=14)
plt.legend(fontsize=12)
plt.grid(True, linestyle='--', alpha=0.7)

# 保存图像
comparison_image = os.path.join(script_dir, 'parameter_comparison2.png')
plt.savefig(comparison_image, dpi=300, bbox_inches='tight')
print(f"参数对比图已保存到: {comparison_image}")
plt.show()

# 2. 绘制损失曲线对比
plt.figure(figsize=(14, 8))
for result in all_results:
    # 训练损失：使用组合颜色 + 实线
    plt.plot(result['history'].history['loss'],
             color=result['color'],
             linestyle='-',
             linewidth=2,
             label=f"{result['name']} (训练)")

    # 验证损失：使用组合颜色 + 虚线
    plt.plot(result['history'].history['val_loss'],
             color=result['color'],
             linestyle='--',
             linewidth=1.5,
             label=f"{result['name']} (验证)")

plt.title('不同参数组合的训练损失对比', fontsize=16)
plt.xlabel('迭代次数', fontsize=12)
plt.ylabel('损失值', fontsize=12)
plt.legend(fontsize=10, ncol=2)  # 分两列显示图例
plt.grid(True, linestyle='--', alpha=0.7)
plt.yscale('log')  # 对数尺度更易观察差异

# 保存图像
loss_comparison_image = os.path.join(script_dir, 'loss_comparison2.png')
plt.savefig(loss_comparison_image, dpi=300, bbox_inches='tight')
print(f"损失对比图已保存到: {loss_comparison_image}")
plt.show()

# ... 后面的代码保持不变 ...
# 3. 输出评估指标表格
print("\n参数组合性能汇总:")
print(
    f"{'名称':<15} | {'窗口':<6} | {'神经元':<10} | {'Dropout':<8} | {'RMSE':<10} | {'MAE':<10} | {'训练时间(秒)':<12}")
print("-" * 90)

for result in all_results:
    print(f"{result['name']:<15} | "
          f"{result['window']:<6} | "
          f"{result['params']['neurons1']}/{result['params']['neurons2']:<10} | "
          f"{result['params']['dropout']:<8.2f} | "
          f"{result['rmse']:<10.4f} | "
          f"{result['mae']:<10.4f} | "
          f"{result['training_time']:<12.2f}")

# 保存结果到CSV
results_df = pd.DataFrame([{
    'Name': r['name'],
    'Window': r['window'],
    'Neurons1': r['params']['neurons1'],
    'Neurons2': r['params']['neurons2'],
    'Dropout': r['params']['dropout'],
    'BatchSize': r['params']['batch_size'],
    'Epochs': r['params']['epochs'],
    'RMSE': r['rmse'],
    'MAE': r['mae'],
    'TrainingTime': r['training_time']
} for r in all_results])

csv_path = os.path.join(script_dir, 'parameter_results2.csv')
results_df.to_csv(csv_path, index=False)
print(f"\n详细结果已保存到: {csv_path}")

# 在脚本末尾添加文件汇总信息
print("\n" + "=" * 50)
print("生成文件汇总:")
print(f"1. 参数对比图: {comparison_image}")
print(f"2. 损失对比图: {loss_comparison_image}")
print(f"3. 详细结果表: {csv_path}")
print("=" * 50)

# 尝试自动打开文件夹
try:
    if os.name == 'nt':  # Windows
        os.startfile(script_dir)
    elif os.name == 'posix':  # Mac/Linux
        subprocess.call(['open' if sys.platform == 'darwin' else 'xdg-open', script_dir])
except Exception as e:
    print(f"无法自动打开文件夹: {e}")
    print(f"请手动访问目录: {script_dir}")